package org.nhindirect.stagent;

import java.io.BufferedInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.util.ArrayList;

import java.util.Enumeration;
import java.util.Iterator;
import java.util.List;

import javax.mail.internet.ContentType;
import javax.mail.internet.MimeMultipart;
import javax.mail.internet.ParseException;
import javax.mail.util.ByteArrayDataSource;

import org.apache.commons.codec.binary.Base64;
import org.nhindirect.stagent.ProtocolException.ProtocolError;

public class SignedEntity extends MultipartEntity
{
    MimeEntity m_content;
    MimeEntity m_signature;
    MimeMultipart m_originalMMPart;
    
    public SignedEntity(ContentType mmContentType, MimeMultipart mm)
    {
    	super(mmContentType);
    	
        if (mm == null)
        {
            throw new IllegalArgumentException();
        }

        m_originalMMPart = mm;
        
    	MimeEntity signature = null;
    	MimeEntity content = null;

        
        try
        {
	
        	String mmCountentType = mm.getContentType();
        	
        	
	        for (int i = 0; i < mm.getCount(); ++i)
	        {
	        	
	        	String contentType = mm.getBodyPart(i).getContentType();
	        	
	        	
	        	if (contentType.contains("application/pkcs7-signature") 
	        			|| contentType.contains("application/x-pkcs7-signature"))
	        	{
	        		signature = new MimeEntity();	
	        		
	        	    Enumeration en = mm.getBodyPart(i).getAllHeaders();
	        	    HeaderCollection headers = new HeaderCollection();
	        	    while (en.hasMoreElements())
	        	    {
	        	    	javax.mail.Header hd = (javax.mail.Header)en.nextElement();
	        	    	headers.add(new Header(hd.getName(), hd.getValue()));	        	    	
	        	    }
	        	    signature.setHeaders(headers);	        	    
	        	    
	        		// get the raw bytes
	                InputStream stream = new BufferedInputStream(mm.getBodyPart(i).getInputStream());
	                
	        	    ByteArrayOutputStream bos = new ByteArrayOutputStream();
	        	    byte[] buffer = new byte[2048];
	        	    int len;
	        	    while ((len = stream.read(buffer, 0, 2048)) != -1)
	        	    	bos.write(buffer, 0, len);
	        	    
	        	    bos.flush();
	        	    
	        	    String str = new String(bos.toByteArray());
	        	    
	        	    // encode the body to Base64
	        	    signature.setBody(new Body(Base64.encodeBase64String(bos.toByteArray())));	

	        		
	        	}
	        	else
	        	{
	        		// get the raw bytes
	                InputStream stream = new BufferedInputStream(mm.getBodyPart(i).getInputStream());
	                
	        	    ByteArrayOutputStream bos = new ByteArrayOutputStream();
	        	    byte[] buffer = new byte[2048];
	        	    int len;
	        	    while ((len = stream.read(buffer, 0, 2048)) != -1)
	        	    	bos.write(buffer, 0, len);

	        	    bos.flush();
	        	    
	        	    content = new MimeEntity();
	        	    content.setBody(new Body(new String(bos.toByteArray(), "ASCII")));
	        	    
	        	    Enumeration en = mm.getBodyPart(i).getAllHeaders();
	        	    HeaderCollection headers = new HeaderCollection();
	        	    while (en.hasMoreElements())
	        	    {
	        	    	javax.mail.Header hd = (javax.mail.Header)en.nextElement();
	        	    	
	        	    	headers.add(new Header(hd.getName(), hd.getValue()));
	        	    	
	        	    }
	        	    
	        	    content.setHeaders(headers);
	        	    
	        	}
	        }
        }
	
        catch (Exception e) {}
	
        this.setContent(content);
        this.setSignature(signature);

    }
    public SignedEntity(DigestAlgorithm algorithm, MimeMultipart mm)
    {
    	this(createContentType(algorithm), mm);
    	
    }
    
    public MimeEntity getContent()
    {
            return this.m_content;
    }

    public void setContent(MimeEntity value)
    {
        if (value == null)
        {
            throw new IllegalArgumentException();
        }
        
        this.m_content = value;
    }
    
    public MimeEntity getSignature()
    {
        return this.m_signature;
    }
    
    public void setSignature(MimeEntity value)
    {
        if (value == null)
        {
            throw new IllegalArgumentException();
        }

        //value.verifyContentType(Protocol.SignatureContentType);
        //value.verifyTransferEncoding(Protocol.TransferEncodingBase64);
        this.m_signature = value;
    }
            
    @Override
    public Iterator<MimeEntity> iterator()
    {
    	List<MimeEntity> tempList = new ArrayList<MimeEntity>();
    	
    	tempList.add(this.m_content);
    	tempList.add(this.m_signature);
    	
    	return tempList.iterator();
    }
    
    public MimeMultipart getMimeMultipart()
    {
    	return m_originalMMPart;
    }
    
    
    public static SignedEntity load(MimeEntity source)
    {
        if (source == null)
        {
            throw new IllegalArgumentException();
        }

        if (!source.isMultiPart())
        {
            throw new IllegalStateException();
        }
        

        
        String sourceText = source.getBody().getSourceText().toString();        
        ByteArrayDataSource dataSource = new ByteArrayDataSource(sourceText.getBytes(), source.getContentType().getValue());
        
        MimeMultipart mm = null;
        
        try
        {
        	mm = new MimeMultipart(dataSource);
        	
        	if (mm.getCount() == 0)
        	{
        		throw new ProtocolException(ProtocolError.InvalidMimeEntity);
        	}
        }
        catch (Exception e)
        {
        	throw new ProtocolException(ProtocolError.InvalidMimeEntity, e);
        }
        
        SignedEntity retVal = null;
        try
        {
        	retVal = new SignedEntity(new ContentType(mm.getContentType()), mm);
        }
    	catch (ParseException e)
    	{
        	throw new ProtocolException(ProtocolError.InvalidHeader, e);
    	}

        return retVal;
    }
       
    
    static ContentType createContentType(DigestAlgorithm digestAlgorithm)
    {
    	ContentType contentType = null;
    	try
    	{
    		contentType = new ContentType(Protocol.MultiPartType_Signed);
    	
    		contentType.setParameter("micalg", Protocol.toString(digestAlgorithm));
    	}
    	catch (ParseException e){}
    	
        return contentType;
    }
}
